{
 "cells": [
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets,transforms\n",
    "import matplotlib.pyplot as plt"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "5c3066a00953255a"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "batch_size = 512\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "trainloader = torch.utils.data.DataLoader(datasets.MNIST('data',train=True,download=True,\n",
    "                                                         transform=transforms.Compose([transforms.ToTensor()])),batch_size=batch_size,shuffle=True)\n",
    "testloader = torch.utils.data.DataLoader(datasets.MNIST('data',train=False,download=False,\n",
    "                                                         transform=transforms.Compose([transforms.ToTensor()])),batch_size=batch_size,shuffle=True)\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "cad8bf7cbe541be4"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net,self).__init__()\n",
    "        self.conv1 = nn.Conv2d (1,6,5,padding=2)\n",
    "        self.conv2 = nn.Conv2d (6,16,5)\n",
    "        self.fc1 = nn.Linear(16*5*5,120)\n",
    "        self.fc2 = nn.Linear(120,84)\n",
    "        self.clf = nn.Linear(84,10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # conv1\n",
    "        x = self.conv1(x)\n",
    "        # 激活函数sigmoid()\n",
    "        x = torch.sigmoid(x)\n",
    "        # 平均池化层，kernel=2x2，步长2\n",
    "        x = F.avg_pool2d(x, 2)\n",
    "        # conv2\n",
    "        x = self.conv2(x)\n",
    "        # 激活函数sigmoid()\n",
    "        x = torch.sigmoid(x)\n",
    "        # 平均池化层，2x2，步长2\n",
    "        x = F.avg_pool2d(x, 2)\n",
    "        # 展平，从第1维开始展平\n",
    "        x = x.view(x.size(0), -1)\n",
    "        # 全连接层1\n",
    "        x = self.fc1(x)\n",
    "        # 激活函数sigmoid()\n",
    "        x = torch.sigmoid(x)\n",
    "        # 全连接层2\n",
    "        x = self.fc2(x)\n",
    "        # 激活函数sigmoid()\n",
    "        x = torch.sigmoid(x)\n",
    "        # 分类层\n",
    "        x = self.clf(x)\n",
    "        return x\n",
    "\n",
    "model = Net().to(device)\n",
    "optimizer = optim.Adam(model.parameters(),lr=1e-2)\n",
    "model"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "798fa43b8e78874"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "epochs = 30\n",
    "accs,losses = [],[]\n",
    "for epochs in range(epochs):\n",
    "    for batch_idx,(x,y) in enumerate(trainloader):\n",
    "        x,y = x.to(device),y.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        out = model(x)\n",
    "        loss = F.cross_entropy(out,y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    correct = 0\n",
    "    testloss = 0\n",
    "    with torch.no_grad():\n",
    "        for batch_idx,(x,y) in enumerate(trainloader):\n",
    "            x,y = x.to(device),y.to(device)\n",
    "            out = model(x)\n",
    "            testloss += F.cross_entropy(out,y).item()\n",
    "            pred = out.max(dim=1,keepdim=True)[1]\n",
    "            correct += pred.eq(y.view_as(pred)).sum().item()\n",
    "\n",
    "    acc = correct/len(testloader.dataset)\n",
    "    testloss = testloss/(batch_idx+1)\n",
    "    accs.append(acc)\n",
    "    losses.append(testloss)\n",
    "    print('epoch:{}，loss:{:.4f},acc：{:.4f}'.format(epochs,testloss,acc))\n",
    "    feature1 = F.sigmoid(model.conv1(x))\n",
    "    #feature1 = F.avg_pool2d(feature1,kernel_size=2,stride=2)\n",
    "    feature2 = F.sigmoid(model.conv2(feature1))\n",
    "    #feature2 = F.avg_pool2d(feature2,kernel_size=2,stride=2)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "5864d9e8a9a924b5"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    " n = 5\n",
    "    img = x.detach().cpu().numpy()[:n]\n",
    "    feature_map1 = feature1.detach().cpu().numpy()[:n]\n",
    "    feature_map2 = feature2.detach().cpu().numpy()[:n]\n",
    "\n",
    "    fig,ax = plt.subplots(3,n,figsize=(10,10))\n",
    "    for i in range(n):\n",
    "        ax[0,i].imshow(img[i].sum(0),cmap='gray')\n",
    "        #ax[0,i].axis('off')\n",
    "        ax[1,i].imshow(feature_map1[i].sum(0),cmap='gray')\n",
    "        #ax[1,i].axis('off')\n",
    "        ax[2,i].imshow(feature_map2[i].sum(0), cmap='gray')\n",
    "        #ax[1,i].axis('off')\n",
    "    plt.show()"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "a64f254289c345d7"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
